import os
import sys

import cv2
from osgeo import gdal
import numpy as np

# 这段代码是用于将经过裁剪的预测图片按照一定的规则拼接起来，并最终保存成一张合成的图像。具体的步骤如下：
#  读取tif数据集
def readTif(fileName):
    dataset = gdal.Open(fileName)
    if dataset == None:
        print(fileName + "文件无法打开")
    return dataset


#  保存tif文件函数
def writeTiff(im_data, im_geotrans, im_proj, path):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    if dataset != None:
        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影
    for i in range(im_bands):
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset


def Tif_merge(TifPath, PrePath, CropSize, RepetitionRate):
    dataset_img = readTif(TifPath)
    width = dataset_img.RasterXSize
    height = dataset_img.RasterYSize
    proj = dataset_img.GetProjection()
    geotrans = dataset_img.GetGeoTransform()
    img = dataset_img.ReadAsArray(0, 0, width, height)  # 获取数据

    img = np.zeros_like(img)

    new_name = 1
    #  裁剪图片,重复率为RepetitionRate
    for i in range(int((height - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        for j in range(int((width - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
            predict_img = cv2.imread(f"{PrePath}/{new_name}.png", 0)

            img[int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
            int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize] = predict_img
            #  文件名 + 1
            new_name = new_name + 1

    #  向前裁剪最后一列
    for i in range(int((height - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        predict_img = cv2.imread(f"{PrePath}/{new_name}.png", 0)
        img[int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize, (width - CropSize): width] = predict_img
        #  写图像

        new_name = new_name + 1
    #  向前裁剪最后一行
    for j in range(int((width - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        predict_img = cv2.imread(f"{PrePath}/{new_name}.png", 0)
        img[(height - CropSize): height, int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize] = predict_img

        #  文件名 + 1
        new_name = new_name + 1

    #  裁剪右下角
    predict_img = cv2.imread(f"{PrePath}/{new_name}.png", 0)
    img[(height - CropSize): height, (width - CropSize): width] = predict_img
    new_name = new_name + 1

    cv2.imwrite('merge_pred.png', np.uint8(img)*255)


if __name__ == '__main__':
    Tif_merge(r"./datasets/mask/test_area_groundtruth.tif", r"./img_out_show", 256, 0.25)
